#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011, Annotation Improvement"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"

"""
TO DO: 
Usage: python annotation_improvement.py tabfile.txt top_blast_hits.txt

"""

import sys
import re
from Bio.Blast import NCBIXML

#Check start sites begins
annofile = sys.argv[1]
af = open(annofile, "rU")
#rec = SeqIO.read(sf, "fasta")
lines = af.readlines()
num = len(lines)

blasthits = sys.argv[2]
bf = open(blasthits, "rU")
blines = bf.readlines()

i = 0

cds_list = []
cds_starts = []
cds_stops = []
cds_desc = []
rna_list = []
rna_defs = []
rna_start_list = []
rna_stop_list = []
crispr_starts = [141632, 3742956, 3820468]
crispr_stops = [153410, 3753501, 3828723]

for line in lines:
    l = line.split('\t')
    feat = l[0]
    locus_tag = l[1]
    start = int(l[2])
    stop = int(l[3])
    frame = l[4]
    gene = l[5]
    protein_id = l[6]
    ec_num = l[7]
    note = l[8]
    bound_moiety = l[9]
    desc = l[10].rstrip()

    if feat == "CDS":
        cds_list.append(locus_tag)
        cds_starts.append(start)
        cds_stops.append(stop)
        cds_desc.append(desc)

    elif feat == "rRNA":
        rna_list.append(locus_tag)
        rna_defs.append(desc)
        rna_start_list.append(start)
        rna_stop_list.append(stop)
        #print locus_tag, feat, start, stop, desc

    elif feat == "tRNA":
        rna_list.append(locus_tag)
        rna_defs.append(desc)
        rna_start_list.append(start)
        rna_stop_list.append(stop)
        #print locus_tag, feat, start, stop, desc
    else:
        pass

acc_dict = {}

for line in blines:
    l = line.split('\t')
    ltag = l[0]
    accession = l[1].rstrip()
    acc_dict[ltag] = accession

rna_count = len(rna_start_list)

#print "Total RNAs found:", rna_count

count = len(cds_list)
i = 0

while i < count:
    #Checking for CRISPR regions
    for j, k in enumerate(crispr_starts):
        #print j, k
        if crispr_starts[j] < cds_starts[i] < crispr_stops[j] and crispr_starts[j] < cds_stops[i] < crispr_stops[j]:
            print cds_list[i], cds_starts[i], cds_stops[i], "\t", "Need to be removed! Contained within CRISPR region", crispr_starts[j], crispr_stops[j]

    #Checking for RNA overlaps
    for j, k in enumerate(rna_list):
        if cds_starts[i] < rna_start_list[j] < cds_stops[i] or cds_starts[i] < rna_stop_list[j] < cds_stops[i]:
            print cds_list[i], cds_starts[i], cds_stops[i], "\t", "Need to be removed! Overlaps with RNA coding region", rna_list[j], rna_defs[j], rna_start_list[j], rna_stop_list[j]
        elif cds_starts[i] < rna_start_list[j] < cds_stops[i] and cds_starts[i] < rna_stop_list[j] < cds_stops[i]:
            print cds_list[i], cds_starts[i], cds_stops[i], "\t", "Need to be removed! Contained within RNA coding region", rna_list[j], rna_defs[j], rna_start_list[j], rna_stop_list[j]

    #Checking for frameshifts
    if i == 0 or i == count - 1:
        pass
    else:
        #print acc_dict[cds_list[i]]
        m1 = re.compile("hypothetical protein")
        if m1.match(cds_desc[i]):
            pass
        else:
            if acc_dict.has_key(cds_list[i]) and acc_dict.has_key(cds_list[i + 1]):
                if acc_dict[cds_list[i]] == acc_dict[cds_list[i + 1]]:
                   print cds_list[i], cds_desc[i], "AND", cds_list[i+1], cds_desc[i+1], "May be frameshifted"

    #Checking for potential pseudogenes (assigned definitions but with >20% domain truncations)

    #Try to check "hypothetical proteins" with conserved domains to assign potential definition
#    m2 = re.compile("hypothetical protein")
#    if m2.match(cds_desc[i]):
#       pass
#    else: #m2.match(cds_desc[i]):
#        rps_pfam = open(cds_list[i] + ".pfam.rpsblast.xml")
#        pfam_rec_gen = NCBIXML.parse(rps_pfam)
#        pfam_rec = pfam_rec_gen.next()
#        for alignment in pfam_rec.alignments:
#            for hsp in alignment.hsps:                
#                if hsp.align_length < 0.6 * alignment.length:
#                    print cds_list[i], cds_desc[i], "Pfam alignment length is less than 60% of full length: (", hsp.align_length, "<", 0.6 * alignment.length, ")"

    i += 1


af.close()
bf.close()